Codex/vae pyhealth2 - #1358

_{Reviewed commit 3e0ff51}

AI Review

suggestion: need change
decision: human review

Thank you for coming back to the VAE bounty and porting it to PyHealth 2.0. Adding categorical EHR sequence support, optional label conditioning for images, and three example notebooks is a useful extension. Line coverage for pyhealth/models/vae.py is 88.29%, above the 80% bar, and total repository coverage is 74.07%, above the 60% bar.

A maintainer will also take a look. The change alters the public contract of the existing VAE: hidden_dim changes meaning and default, latent_dim is new, and the generate signature changes. Existing correct calls can therefore break or behave differently, and a maintainer should decide whether that is acceptable.

To get this merged:

1. Single-label image datasets. The constructor auto-uses a sole output label as the condition. A dataset with a single multilabel or regression output, for example ChestXray14 multilabel, now raises an error, where the old version ignored labels. Only auto-condition when the label is binary or multiclass, and otherwise ignore it unless condition_key is given explicitly.
2. generate(num_samples=...) on labeled image datasets. This now raises because of the auto-conditioning. Please keep unconditional generation working, for example by sampling the condition at random or only conditioning when condition_key is passed explicitly.
3. @​torch.no_grad() removed from generate. The returned images now require grad, so .numpy() and similar calls fail. Please restore the decorator or detach the output. The old device argument has also been dropped, and the second positional argument is now sequence_length. Please keep the signature backward compatible or call the change out clearly.
4. hidden_dim semantics. Previously hidden_dim was the latent size, with a default of 128. It is now the channel width, with a default of 32. Existing VAE(dataset, hidden_dim=N) calls silently build a different network. Please keep hidden_dim as the latent size, or document the break prominently in the description and docs.
5. Tests. The COVID19CXRImageGeneration task tests and the learning test were deleted, although that task still ships. Please restore tests for the task and a test that training reduces the reconstruction error. Also add tests for the new error branches and for Trainer.train.

Optional suggestions:

• In the ChestXray14 notebook, the comment about patient-level splits being streaming subsets looks inaccurate. split_by_patient returns subsets that get_dataloader should accept, so you could use it directly.
• Unsupported PIL modes such as "P" raise a bare KeyError at the mode lookup in _init_image. A clear ValueError would be friendlier.

_{Reviewed commit 3e0ff51}

_{Reviewed commit 478feab}

AI Review

suggestion: need change
decision: human review

Thank you for this contribution. It extends the PyHealth 2.0 VAE to categorical EHR code sequences and optional label-conditioned image generation, and it adds three example notebooks and docs updates. The new model code looks reasonable. A maintainer will also take a look, because the rewrite of pyhealth/models/vae.py changes the existing VAE's public interface and the format of saved checkpoints.

Coverage: pyhealth/models/vae.py is at 89.05%, above the 80% bar for modified files, and the repository total is 74.11%, above the 60% bar. I did not run the tests or notebooks.

To get this merged:

1. Keep backward compatibility, or state the break clearly. The old encode() and decode() methods are gone. The attributes encoder, decoder, mu and logvar are renamed, and the ResBlock2D architecture is replaced. State_dicts saved from the previous VAE therefore no longer load. Please restore encode and decode (or add thin wrappers). If the architecture change is intentional, document it in the PR description and the docs as a breaking change.
2. Image processors without an explicit mode used to work, because the channel count was read from the data. They now raise a ValueError. Please infer the channel count from a sample when mode is None, or document this as an intentional break. Also keep a clear error for to_tensor=False, which the old to_tensor check used to catch.
3. Restore the test coverage that was removed. The multi-epoch "learns to reconstruct" test is now a single optimizer step. The RGB, gradient-flow and state_dict round-trip tests are gone, so the RGB path used by the COVID classification example is untested. Please add at least an RGB forward-pass test and a state_dict round-trip test.

Optional suggestions:

• The sequence path always samples the latent, even in eval mode, so validation loss is noisy. Using mu when not self.training would make it deterministic.
• In the synthetic notebook, torch.manual_seed before an eval-mode forward has no effect, so the remark about seeding can be dropped.
• The ChestXray14 notebook subclasses the binary task only to add patient_id. That would be better placed in the task itself.

_{Reviewed commit 478feab}

_{Reviewed commit f2973cd}

_{Reviewed commit bed99d7}

AI Review

suggestion: need change
decision: human review

Thank you for this contribution. It extends the PyHealth 2.0 VAE with optional label conditioning for images and a new categorical EHR sequence VAE. It also adds three example notebooks, documentation updates and tests. It keeps the earlier image VAE's module names and state_dict keys, which is a nice touch.

A maintainer will take a look as well. The change modifies pyhealth/trainer.py, a shared base class, so Trainer.inference() now behaves differently for any model that returns mask and reconstruction_logits. That is a core-level change that needs a human decision.

To get this merged:

1. In VAE.generate() for sequence VAEs, the latent is sampled on the device argument. The latent_to_hidden, sequence_decoder and sequence_output layers live on the model's device. If the two differ, generate(device=...) will fail. Either move the latent to self.device, or document that device must match the model's device. Please add a test for this path.

Coverage: vae.py is at 87.0% and chestxray14_binary_classification.py is at 100%, both above the 80% bar. trainer.py is at 77.2%, just under 80%, and the whole repository is at 75.6%, above the 60% bar. The trainer.py shortfall mostly comes from existing code. Still, the new variable-length branch in inference() is tested, so it would be good to confirm that it is covered.

Optional suggestions:

• The image forward pass no longer accepts a list of image tensors, as the previous version did. Consider restoring that or mentioning the change in the docs.
• Trainer.inference() stores full per-token probability arrays for sequence VAEs. With large EHR vocabularies this could use a lot of memory, so a note in the docs may help.
• In the ChestXray14 notebook, the elif not DATA_AVAILABLE branch is unreachable, because DOWNLOAD_DATA is already not DATA_AVAILABLE. The notebook also starts a roughly 5k-image download without asking, so a clearer warning may help.
• Conditioning is tested only for output shape. A test showing that different conditions change the output would strengthen it.
• The synthetic notebook describes training results but has no saved outputs.

_{Reviewed commit bed99d7}